#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Han Xiao <artex.xh@gmail.com> <https://hanxiao.github.io>

# NOTE: First install bert-as-service via
# $
# $ pip install bert-serving-server
# $ pip install bert-serving-client
# $

# simple similarity search on FAQ

import numpy as np
from bert_serving.client import BertClient
from termcolor import colored
import random
# prefix_q = '##### **Q:** '
# prefix_a = '**A:** '
topk = 5
def next_print(flag,s):
    try:
        a = 0
        for idx in topk_idx:
            a += 1
            if a == 2:
                flag = True
            print('> %s\t%s' % (colored('%.1f' % score[idx], 'cyan'), colored(answers[idx], 'yellow')))
    except:
        if flag:
            pass
        else:
            print(colored(s, "red"))


with open('s.txt',encoding="utf-8") as fp:
    questions = []
    for i in fp.readlines():cons = i.split(",");q = "".join([v for i,v in enumerate(list(cons[0])) if (i != 0 and i != len(list(cons[0])) - 1) and v !=""]);questions.append(q)
    questions = [i for i in questions if i != ""]
    print('%d questions loaded, avg. len of %d' % (len(questions), np.mean([len(d.split()) for d in questions])))
with open('s.txt',encoding="utf-8") as fp:
    answers = []
    for i in fp.readlines(): cons = i.split(",");q = "".join([v for i, v in enumerate(list(cons[1])) if (i != 0 and i != len(list(cons[1])) - 1) and v != ""]);answers.append(q)
print('%d answers loaded, avg. len of %d' % (len(answers), np.mean([len(d.split()) for d in answers])))

error_list = ["这不在我的范围内","请问一些关于国家的问题","我还在发育哈","其他的都可以呦"]
with BertClient() as bc:
    doc_vecs = bc.encode(questions)

    while True:
        query = input(colored('your question: ', 'green'))
        if query == "":
            print("输入不能为空！")
            continue
        query_vec = bc.encode([query])[0]
        # compute normalized dot product as score
        score = np.sum(query_vec * doc_vecs, axis=1) / np.linalg.norm(doc_vecs, axis=1)
        topk_idx = np.argsort(score)[::-1][:topk]
        print('top %d questions similar to "%s"' % (topk, colored(query, 'green')))
        flag = False
        s = random.choice(error_list)
        next_print(flag,s)